package com.cyf.common.core.thread;

import java.io.Closeable;
import java.io.IOException;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.*;

/**
 *  线程同步结束器<br>
 *  允许一个或多个线程在执行任务时，一直等待，保证一起运行
 *
 *  实现 Closeable 接口，可以使用 try(SyncFinisher ...) 自动关闭资源
 */
public class SyncFinisher implements Closeable {

    private final Set<Worker> workers;
    private final int threadSize;

    /**
     * 默认多个线程同时开始
     */
    private boolean isBeginAtSameTime = true;

    private ExecutorService executorService;

    /**
     * 启动同步器，用于保证所有worker线程同时开始
     */
    private final CountDownLatch beginLatch;

    /**
     * 结束同步器，用于等待所有worker线程同时结束，可选
     */
    private CountDownLatch endLatch;

    public SyncFinisher(int threadSize) {
        this.workers = new HashSet<>();
        this.threadSize = threadSize;
        this.beginLatch = new CountDownLatch(1);
    }

    /**
     * 设置是否所有 work 线程同时开始
     * @param isBeginAtSameTime
     * @return
     */
    public void setBeginAtSameTime(boolean isBeginAtSameTime) {
        this.isBeginAtSameTime = isBeginAtSameTime;
    }

    /**
     * 自定义线程池
     * @param corePoolSize
     * @return
     */
    private ThreadPoolExecutor newExecutor(int corePoolSize) {
        return new ThreadPoolExecutor(
                corePoolSize, // 初始线程数为 corePoolSize 指定的大小
                Integer.MAX_VALUE, // 没有最大线程数限制
                60, TimeUnit.SECONDS,
                new LinkedBlockingQueue<>(),  // 默认使用LinkedBlockingQueue，默认队列大小为1024
                Executors.defaultThreadFactory(),
                new ThreadPoolExecutor.AbortPolicy()); // 如果队列满了，最新的提交任务将会被拒绝
    }


    /**
     * 增加工作线程
     * @param task
     * @return
     */
    public void addWorker(Runnable task) {
        Worker worker = new Worker() {
            @Override
            public void work() {
                task.run();
            }
        };
        workers.add(worker);
    }

    /**
     *  开始工作
     * @param sync 主线程是否阻塞等待
     */
    public void start(boolean sync) {
        endLatch = new CountDownLatch(threadSize);

        if (executorService == null || executorService.isShutdown()) {
            this.executorService = newExecutor(threadSize);
        }
        for (Worker worker : workers) {
            executorService.submit(worker);
        }
        // 保证所有 worker 同时开始
        beginLatch.countDown();

        if(sync) {
            try {
                endLatch.await();
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 剩余任务数
     * @return
     */
    public long count() {
        return endLatch.getCount();
    }

    /**
     * 结束线程池
     */
    public void stop() {
        if (executorService != null) {
            this.executorService.shutdown();
        }
        this.executorService = null;
        workers.clear();
    }

    @Override
    public void close() throws IOException {
        stop();
    }

    public abstract class Worker implements Runnable {

        @Override
        public void run() {
            // 如果设置线程同时开始执行任务，则等待主线程将任务都提交给线程池，否则直接运行
            if (isBeginAtSameTime) {
                try {
                    beginLatch.await();
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
            try {
                work(); // work 即使出现异常，保证结束同步器减 1
            } finally {
                endLatch.countDown();
            }
        }

        public abstract void work();

    }
}
